from QKDNetwork import QKDNetwork
from docplex.mp.model import Model

def getPSPHETimeSlot(routeList: list, net: QKDNetwork):
    # routeList: [{"sd": [1, 2, 4], "path": [1, 3, 2]}...]
    model = Model(name="PSP_HE_Time_Slot")
    T = model.continuous_var(name="T", lb=0)
    M = 1000
    
    p_dim1, p_dim2 = net.num_nodes + 1, net.num_nodes + 1
    p = model.integer_var_dict(
        [(i, j) for i in range(p_dim1) for j in range(p_dim2)],
        lb=0,
        name="p"
    )
    
    def is_sublist(sub, main):
        sub_length = len(sub)
        return any(sub == main[i:i+sub_length] for i in range(len(main) - sub_length + 1))
    
    for u in range(net.num_nodes):
        for v in net.G.neighbors(u):
            flow = model.linear_expr()
            flow += p[u, v]
            flow += p[v, u]
            for i, item in enumerate(routeList):
                path = item["path"]
                if(is_sublist([u, v], path) or is_sublist([v, u], path)):
                    flow -= 1
            model.add_constraint(flow >= 0)
            model.add_constraint(p[u, v] <= M * net.G.nodes[u]["transmitter"])
        
    for i in net.mdi_links:
        model.add_constraint(p[i[0], i[1]] == p[i[2], i[1]])
    
    for u in range(net.num_nodes):
        sum_p = model.linear_expr()
        for v in net.G.neighbors(u):
            sum_p += p[u, v]
        model.add_constraint(sum_p <= T)
    
    model.minimize(T)
    sol = model.solve()
    if sol:
        return T.solution_value
    
    return 0
    